import chisel3._
import chisel3.util._

/**
  * 指令Cache模块
  * @param SIZE      Cache行数
  * @param LINE_SIZE 每行字节数
  * @param WAY_NUM   组相连路数
  * @param UseLRU    是否使用LRU替换策略
  */
class MyICache(val SIZE: Int, val LINE_SIZE: Int, val WAY_NUM: Int, val UseLRU: Boolean)
    extends Module {
  // 各项参数宽度定义
  val INDEX_WIDTH   = log2Ceil(SIZE)
  val OFFSET_WIDTH  = log2Ceil(LINE_SIZE)
  val WAY_NUM_WIDTH = log2Ceil(WAY_NUM)
  val LINE_BIT_SIZE = LINE_SIZE * 8                      // 每行为多少bit
  val LINE_WORD_NUM = LINE_SIZE / 4                      // 每行为多少word

  // 打印cache主要配置信息
  println("Inst Cache Size: " + SIZE)
  println("Inst Cache Line Size   : " + LINE_WORD_NUM + " Word | " + LINE_SIZE + " Byte | " + LINE_BIT_SIZE + " Bit")
  println("Inst Cache Way Number  : " + WAY_NUM)

  // 检查索引宽度与offset宽度的和不得超过12
  assert(
    (INDEX_WIDTH + OFFSET_WIDTH <= 12).asBool,
    "Index windth + offset width should be less than or equal 12, but given is %d\n",
    (INDEX_WIDTH + OFFSET_WIDTH).asUInt
  )

  // IO定义
  val io = IO(new Bundle {
    // cpu接口
    val valid     = Input(Bool())           // 请求有效
    val op        = Input(Bool())           // 读写操作（假设，具体视设计而定）
    val tag       = Input(UInt(20.W))       // 地址tag
    val index     = Input(UInt(INDEX_WIDTH.W)) // 行索引
    val offset    = Input(UInt(OFFSET_WIDTH.W)) // 行内偏移
    val addr_ok   = Output(Bool())          // 地址握手信号
    val data_ok   = Output(Bool())          // 返回数据握手信号
    val rdata_l   = Output(UInt(32.W))      // 低32位数据
    val rdata_h   = Output(UInt(32.W))      // 高32位数据
    val uncached  = Input(Bool())           // 非cache访问请求
    val cacop_en         = Input(Bool())    // cacop操作请求
    val cacop_op         = Input(UInt(2.W)) // cacop操作类型
    val icache_unbussy   = Output(Bool())   // cache是否空闲
    val tlb_excp_cancel_req   = Input(Bool()) // TLB异常取消请求

    // AXI接口
    val rd_req      = Output(Bool())           // 读请求
    val rd_type     = Output(UInt(3.W))        // 读类型
    val rd_addr     = Output(UInt(32.W))       // 读地址
    val rd_rdy      = Input(Bool())            // 读就绪
    val ret_valid   = Input(Bool())            // AXI返回数据有效
    val ret_last    = Input(Bool())            // AXI读返回数据结束
    val ret_data    = Input(UInt(32.W))        // AXI返回数据
    val wr_req      = Output(Bool())           // 写请求
    val wr_type     = Output(UInt(3.W))        // 写类型
    val wr_addr     = Output(UInt(32.W))       // 写地址
    val wr_wstrb    = Output(UInt(4.W))        // 写选通信号
    val wr_data     = Output(UInt(LINE_BIT_SIZE.W)) // 写数据
    val wr_rdy      = Input(Bool())            // 写就绪
  })

  // 请求/未命中统计器，用于性能统计
  val total_req_counter   = RegInit(0.U(32.W))
  val total_miss_counter  = RegInit(0.U(32.W))
  val dcache_counter      = RegInit(0.U(32.W))
  dcache_counter := dcache_counter + 1.U

  // cache数据/标记/有效bit存储结构，每一路一个module
  val DataBank                = Array.fill(WAY_NUM)(Module(new DataRAM(SIZE, LINE_BIT_SIZE)).io)
  val TagvBank                = Array.fill(WAY_NUM)(Module(new TagvRAM(SIZE, 20)).io)
  val Valid                   = RegInit(VecInit(Seq.fill(SIZE)(0.U(WAY_NUM.W)))) // 有效位

  // 命中信息
  val way_hit                 = Wire(Vec(WAY_NUM, Bool()))      // 每一路是否命中
  val cache_hit               = Wire(Bool())                    // 是否命中任意一路
  val cache_hit_way           = Wire(UInt(WAY_NUM_WIDTH.W))     // 命中的way编号

  // 被替换路缓存，用于替换选路
  val replaced_way_buffer     = RegInit(0.U(WAY_NUM_WIDTH.W))
  val replaced_way            = Wire(UInt(WAY_NUM_WIDTH.W))

  // AXI需要读/写的标志
  val need_write_axi      = RegInit(false.B)
  val need_read_axi       = RegInit(false.B)

  // 每一路读出的数据
  val read_data           = Wire(Vec(WAY_NUM, UInt(LINE_BIT_SIZE.W)))
  val target_data         = Wire(Vec(LINE_WORD_NUM, UInt(32.W))) // 目标行分割为word
  val read_tagv           = Wire(Vec(WAY_NUM, UInt(21.W)))
  val read_tag            = Wire(Vec(WAY_NUM, UInt(20.W)))
  val read_valid          = Wire(Vec(WAY_NUM, Bool()))

  // 请求相关buffer和信号
  val receive_request     = Wire(Bool())
  val uncached_write      = Wire(Bool())
  val uncached_read       = Wire(Bool())
  val uncached_en         = Wire(Bool())
  val read_hit            = Wire(Bool())
  val write_hit           = Wire(Bool())
  val cache_unbussy       = Wire(Bool())

  // 状态机定义，7个主要状态
  object CacheState extends ChiselEnum {
    val IDLE, LOOKUP, WRITEHIT, MISS, REPLACE, REFILL, RESPOND = Value
    val correct_annotation_map = Map[String, UInt](
      "IDLE" -> 0.U,
      "LOOKUP" -> 1.U,
      "WRITEHIT" -> 2.U,
      "MISS" -> 3.U,
      "REPLACE" -> 4.U,
      "REFILL" -> 5.U,
      "RESPOND" -> 6.U
    )
  }
  val main_state  = RegInit(CacheState.IDLE)

  // 各状态专用wire，便于状态判断
  val is_idle     = WireInit(main_state === CacheState.IDLE)
  val is_lookup   = WireInit(main_state === CacheState.LOOKUP)
  val is_writehit = WireInit(main_state === CacheState.WRITEHIT)
  val is_miss     = WireInit(main_state === CacheState.MISS)
  val is_replace  = WireInit(main_state === CacheState.REPLACE)
  val is_refill   = WireInit(main_state === CacheState.REFILL)
  val is_respond  = WireInit(main_state === CacheState.RESPOND)
  dontTouch(is_refill)

  // 其他各种buffer
  val rd_type_buffer    = RegInit(0.U(3.W))
  val rd_addr_buffer    = RegInit(0.U(32.W))
  val wr_type_buffer    = RegInit(0.U(3.W))
  val wr_addr_buffer    = RegInit(0.U(32.W))
  val wr_wstrb_buffer   = RegInit(0.U(4.W))
  val wr_data_buffer    = RegInit(0.U(LINE_BIT_SIZE.W))

  val cache_hit_buffer          = RegInit(false.B)
  val request_cacop_en_buffer   = RegInit(false.B)
  val request_cacop_op_buffer   = RegInit(0.U(2.W))
  val request_op_buffer         = RegInit(false.B)
  val request_tag_buffer        = RegInit(0.U(20.W))
  val request_index_buffer      = RegInit(0.U(INDEX_WIDTH.W))
  val request_offset_buffer     = RegInit(0.U(OFFSET_WIDTH.W))
  val request_uncached_buffer   = RegInit(false.B)
  val target_word_id            = RegInit(0.U((OFFSET_WIDTH-2).W))
  val request_addr              = Cat(Mux(is_lookup, io.tag, request_tag_buffer), request_index_buffer, request_offset_buffer)
  val request_line_addr         = Cat(Mux(is_lookup, io.tag, request_tag_buffer), request_index_buffer, 0.U(OFFSET_WIDTH.W))
  val normal_request            = Wire(Bool())

  val ret_data_buffer   = RegInit(VecInit.fill(LINE_WORD_NUM)(0.U(32.W))) // 装填返回数据用于AXI返回分段
  val miss_data_ret_cnt = RegInit(0.U((OFFSET_WIDTH - 2).W)) // AXI读取返回计数

  // cacop指令解析
  val cacop_op_0  = request_cacop_en_buffer && request_cacop_op_buffer === 0.U
  val cacop_op_1  = request_cacop_en_buffer && request_cacop_op_buffer === 1.U
  val cacop_op_2  = request_cacop_en_buffer && request_cacop_op_buffer === 2.U
  val cacop_way   = request_addr(WAY_NUM_WIDTH-1, 0)

  // 被替换的way选择逻辑，多条件选择
  replaced_way    := Mux(cacop_op_2 && cache_hit, cache_hit_way, 
                        Mux(cacop_op_0 || cacop_op_1, cacop_way, dcache_counter(WAY_NUM_WIDTH-1, 0)))   

  // 将tagv拉开成tag和valid信号
  for (i <- 0 until WAY_NUM) {
    read_tag(i)   := read_tagv(i)(20, 1)
    read_valid(i) := read_tagv(i)(0)
  }
  // 命中判定
  for (i <- 0 until WAY_NUM) {
    way_hit(i) := (read_tag(i) === io.tag) && read_valid(i)
  }

  // 命中数据提取
  target_data     := read_data(cache_hit_way).asTypeOf(target_data)
  cache_hit       := way_hit.reduce(_||_) && !io.tlb_excp_cancel_req
  cache_hit_way   := OHToUInt(way_hit)
  uncached_en     := Mux(is_lookup, io.uncached, request_uncached_buffer) && !request_cacop_en_buffer 
  uncached_read   := uncached_en && !request_op_buffer
  uncached_write  := uncached_en && request_op_buffer
  normal_request  := !Mux(is_lookup, io.uncached, request_uncached_buffer) && !request_cacop_en_buffer
  read_hit        := cache_hit && !request_op_buffer && normal_request
  write_hit       := cache_hit && request_op_buffer && normal_request
  cache_unbussy   := is_idle || is_lookup && cache_hit && normal_request || is_respond
  receive_request := cache_unbussy && (io.valid || io.cacop_en)
  io.icache_unbussy:= cache_unbussy
  io.addr_ok      := cache_unbussy
  io.data_ok      := is_lookup && (cache_hit || request_op_buffer) && !io.uncached || is_respond
  io.rdata_l      := Mux(is_respond, ret_data_buffer(target_word_id), target_data(target_word_id))
  io.rdata_h      := Mux(is_respond, ret_data_buffer(target_word_id+1.U), target_data(target_word_id+1.U))

  io.rd_req       := is_miss && need_read_axi 
  io.rd_type      := rd_type_buffer
  io.rd_addr      := rd_addr_buffer

  io.wr_req       := false.B
  io.wr_type      := 0.U
  io.wr_wstrb     := 0.U
  io.wr_addr      := 0.U
  io.wr_data      := 0.U

  // 数据Bank写入时写冲突处理
  val conflcit      = Wire(Vec(WAY_NUM, Bool()))
  val conflcit_flag = RegInit(VecInit.fill(WAY_NUM)(false.B))
  val conflcit_data = RegInit(VecInit.fill(WAY_NUM)(0.U(LINE_BIT_SIZE.W)))
  for (i <- 0 until WAY_NUM) {
    read_data(i)        := Mux(conflcit_flag(i), conflcit_data(i), DataBank(i).doutb)
    DataBank(i).clka    := clock
    DataBank(i).clkb    := clock
    DataBank(i).addra   := request_index_buffer
    DataBank(i).addrb   := Mux(conflcit(i), io.index + 1.U, io.index) // 有冲突调整读址
    DataBank(i).dina    := ret_data_buffer.asUInt
    DataBank(i).wea     := is_respond && normal_request && (i.asUInt === replaced_way_buffer)  
    conflcit(i)         := request_index_buffer === io.index && DataBank(i).wea
    when(conflcit(i))
    {
      conflcit_flag(i) := true.B;
      conflcit_data(i) := DataBank(i).dina;
    }
    .otherwise
    {
      conflcit_flag(i) := false.B;
      conflcit_data(i) := 0.U(LINE_BIT_SIZE.W);
    }                  
  }
  // 有效位写入和Tagv读写
  val valid_in = Wire(Vec(WAY_NUM, UInt(1.W)))
  for (i <- 0 until WAY_NUM) {
    read_tagv(i)        := Cat(TagvBank(i).douta, Valid(request_index_buffer)(i))
    TagvBank(i).clka    := clock
    TagvBank(i).addra   := Mux(receive_request, io.index, request_index_buffer)
    TagvBank(i).dina    := Mux(request_cacop_en_buffer, 0.U(20.W), request_tag_buffer)
    valid_in(i)         := Mux(TagvBank(i).wea, !request_cacop_en_buffer, Valid(request_index_buffer)(i))
    TagvBank(i).wea     := is_miss && !uncached_en && ((i.asUInt === replaced_way_buffer) || request_cacop_en_buffer)
  }
  Valid(request_index_buffer) := valid_in.asUInt

  // 主状态机
  when(is_idle) { // 空闲，等新请求
    when(receive_request) 
    {
      // 记录请求信息
      request_cacop_en_buffer := io.cacop_en
      request_cacop_op_buffer := io.cacop_op
      request_op_buffer       := io.op
      request_index_buffer    := io.index
      request_offset_buffer   := io.offset
      target_word_id          := io.offset(OFFSET_WIDTH-1, 2)
      main_state := CacheState.LOOKUP
    }
  }
  .elsewhen(is_lookup) 
  {
    // 查询命中情况
    total_req_counter := total_req_counter + 1.U
    request_uncached_buffer   := io.uncached
    request_tag_buffer        := io.tag
    cache_hit_buffer          := cache_hit
    when(io.tlb_excp_cancel_req) 
    {
      main_state := CacheState.IDLE
    }
    .elsewhen(!cache_hit || request_cacop_en_buffer || uncached_en)
    { 
      when(uncached_en) { target_word_id := 0.U }
      replaced_way_buffer   := replaced_way
      need_read_axi       := uncached_read || normal_request
      rd_type_buffer      := Mux(uncached_en, "b010".U, "b100".U)
      rd_addr_buffer      := Mux(uncached_en, request_addr, request_line_addr)
      main_state := Mux(cacop_op_2 && !cache_hit, CacheState.IDLE, CacheState.MISS) 
      total_miss_counter := total_miss_counter + 1.U
    }
    .elsewhen(receive_request) { // 新请求直通LOOKUP
      request_cacop_en_buffer := io.cacop_en
      request_cacop_op_buffer := io.cacop_op
      request_op_buffer       := io.op
      request_index_buffer    := io.index
      request_offset_buffer   := io.offset
      target_word_id          := io.offset(OFFSET_WIDTH-1, 2)
      main_state := CacheState.LOOKUP
    }
    .otherwise { main_state := CacheState.IDLE }
  }
  .elsewhen(is_miss) 
  { 
    // 发出AXI读命令，等待Ready
    when(need_read_axi)
    {
      when(io.rd_rdy)
      { 
        main_state := CacheState.REFILL
        miss_data_ret_cnt := 0.U
        ret_data_buffer   := (0.U).asTypeOf(ret_data_buffer)
      }
    }
    .otherwise { main_state := CacheState.IDLE }
  }
  .elsewhen(is_refill) 
  {
    // AXI分片数据填充Buffer
    when(io.ret_valid) 
    {
      miss_data_ret_cnt := miss_data_ret_cnt + 1.U
      ret_data_buffer(miss_data_ret_cnt) := io.ret_data
      when(io.ret_last) { main_state := CacheState.RESPOND }
    }
  }
  .elsewhen(is_respond) { // 响应数据  
    when(receive_request)
    {
      request_cacop_en_buffer := io.cacop_en
      request_cacop_op_buffer := io.cacop_op
      request_op_buffer       := io.op
      request_index_buffer    := io.index
      request_offset_buffer   := io.offset
      target_word_id          := io.offset(OFFSET_WIDTH-1, 2)
      main_state := CacheState.LOOKUP
    }
    .otherwise {main_state := CacheState.IDLE} 
  }

  // 打印miss率（可选）
  // when(dcache_counter(15,0).asUInt===0.U)
  // {
  //   printf(cf"[DATA] cache miss rate: $total_miss_counter%d $total_req_counter%d\n")
  // }
}